# utils.py

import numpy as np
import torch
from config import INPUT_DIM_FREQ, INPUT_DIM_BAND

def get_state_vector(state):
    # 将环境状态转换为神经网络输入
    threats = state["threats"]
    segments = state["segments"]

    # 构造干扰频段决策器输入
    freq_input = np.array(threats + [len(segments)] + [s[0] for s in segments])
    freq_input = torch.tensor(freq_input, dtype=torch.float32).unsqueeze(0)

    # 构造干扰带宽决策器输入
    band_input = np.array([len(segments)] + [len(s) for s in segments])
    band_input = torch.tensor(band_input, dtype=torch.float32).unsqueeze(0)

    return {
        "freq": freq_input,
        "band": band_input
    }